{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5ae61abe-1f74-4cb4-a0a0-3e0a6b4fe1b9",
   "metadata": {},
   "source": [
    "## 【实操】使用Milvus构建RAG的案例"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40c823d2-a98b-42d2-901c-d117feb2043b",
   "metadata": {},
   "source": [
    "### 准备工作"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d2d2800-f238-43f6-acd7-0023f9222053",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 安装依赖和环境pip install pymilvus ollama\n",
    "# 准备数据\n",
    "# 下载并解压https://github.com/milvus-io/milvus-docs/releases/download/v2.4.6-preview/milvus_docs_2.4.x_en.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "62ae8b8f-eff1-4724-b0e2-735ca5b6d434",
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "\n",
    "text_lines = []\n",
    "\n",
    "for file_path in glob(\"/en/faq/*.md\", recursive=True):\n",
    "    with open(file_path, \"r\") as file:\n",
    "        file_text = file.read()\n",
    "\n",
    "    text_lines += file_text.split(\"# \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f13f9072-3fee-42d5-ac3d-815d7659e061",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ollama\n",
    "\n",
    "def emb_text(text):\n",
    "    response = ollama.embeddings(model=\"mxbai-embed-large\", prompt=text)\n",
    "    return response[\"embedding\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f900eba1-8023-4b8a-b9c9-2ebb453542af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1024\n",
      "[0.23273247480392456, 0.4258939325809479, 0.19729609787464142, 0.46137019991874695, -0.4603538513183594, -0.14180253446102142, -0.18286575376987457, -0.07608925551176071, 0.39981168508529663, 0.8336679935455322]\n"
     ]
    }
   ],
   "source": [
    "test_embedding = emb_text(\"This is a test\")\n",
    "embedding_dim = len(test_embedding)\n",
    "print(embedding_dim)\n",
    "print(test_embedding[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66de23c6-7535-4dce-ab53-ff8360877fd8",
   "metadata": {},
   "source": [
    "### 将数据加载到 Milvus 中"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d7d8229c-d519-4e8e-87f0-17c2152e3f53",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import MilvusClient\n",
    "\n",
    "milvus_client = MilvusClient(uri=\"./milvus_demo.db\")\n",
    "\n",
    "collection_name = \"my_rag_collection\"\n",
    "\n",
    "# 检查 Collections 是否已存在，如果已存在，则删除它\n",
    "if milvus_client.has_collection(collection_name):\n",
    "    milvus_client.drop_collection(collection_name)\n",
    "  \n",
    "milvus_client.create_collection(\n",
    "    collection_name=collection_name,\n",
    "    dimension=embedding_dim,\n",
    "    metric_type=\"IP\",  # Inner product distance\n",
    "    consistency_level=\"Strong\"\n",
    ")\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ae149c45-7b11-4e38-9238-7d7af004fa06",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Creating embeddings: 100%|██████████████████████| 72/72 [00:39<00:00,  1.84it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'insert_count': 72, 'ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71], 'cost': 0}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 插入数据\n",
    "from tqdm import tqdm\n",
    "\n",
    "data = []\n",
    "\n",
    "for i, line in enumerate(tqdm(text_lines, desc=\"Creating embeddings\")):\n",
    "    data.append({\"id\": i, \"vector\": emb_text(line), \"text\": line})\n",
    "\n",
    "milvus_client.insert(collection_name=collection_name, data=data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "337f9ee1-0103-48e1-ae29-b23e9d794a71",
   "metadata": {},
   "source": [
    "### 构建RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "246ff29a-15e1-4c2e-a750-8e533a39405c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 问题\n",
    "question = \"How is data stored in milvus?\"\n",
    "\n",
    "# 在 Collections 中搜索该问题并检索语义前 3 个匹配项\n",
    "search_res = milvus_client.search(\n",
    "    collection_name=collection_name,\n",
    "    data=[\n",
    "        emb_text(question)\n",
    "    ],  # Use the `emb_text` function to convert the question to an embedding vector\n",
    "    limit=3,  # Return top 3 results\n",
    "    search_params={\"metric_type\": \"IP\", \"params\": {}},  # Inner product distance\n",
    "    output_fields=[\"text\"],  # Return the text field\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "10e0eec3-baf2-4070-805f-17aac2caf59f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "    [\n",
      "        \" Where does Milvus store data?\\n\\nMilvus deals with two types of data, inserted data and metadata. \\n\\nInserted data, including vector data, scalar data, and collection-specific schema, are stored in persistent storage as incremental log. Milvus supports multiple object storage backends, including [MinIO](https://min.io/), [AWS S3](https://aws.amazon.com/s3/?nc1=h_ls), [Google Cloud Storage](https://cloud.google.com/storage?hl=en#object-storage-for-companies-of-all-sizes) (GCS), [Azure Blob Storage](https://azure.microsoft.com/en-us/products/storage/blobs), [Alibaba Cloud OSS](https://www.alibabacloud.com/product/object-storage-service), and [Tencent Cloud Object Storage](https://www.tencentcloud.com/products/cos) (COS).\\n\\nMetadata are generated within Milvus. Each Milvus module has its own metadata that are stored in etcd.\\n\\n###\",\n",
      "        231.93528747558594\n",
      "    ],\n",
      "    [\n",
      "        \"How does Milvus flush data?\\n\\nMilvus returns success when inserted data are loaded to the message queue. However, the data are not yet flushed to the disk. Then Milvus' data node writes the data in the message queue to persistent storage as incremental logs. If `flush()` is called, the data node is forced to write all data in the message queue to persistent storage immediately.\\n\\n###\",\n",
      "        226.47874450683594\n",
      "    ],\n",
      "    [\n",
      "        \"What is the maximum dataset size Milvus can handle?\\n\\n  \\nTheoretically, the maximum dataset size Milvus can handle is determined by the hardware it is run on, specifically system memory and storage:\\n\\n- Milvus loads all specified collections and partitions into memory before running queries. Therefore, memory size determines the maximum amount of data Milvus can query.\\n- When new entities and and collection-related schema (currently only MinIO is supported for data persistence) are added to Milvus, system storage determines the maximum allowable size of inserted data.\\n\\n###\",\n",
      "        210.61297607421875\n",
      "    ]\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "# 查看检索结果\n",
    "import json\n",
    "\n",
    "retrieved_lines_with_distances = [\n",
    "    (res[\"entity\"][\"text\"], res[\"distance\"]) for res in search_res[0]\n",
    "]\n",
    "print(json.dumps(retrieved_lines_with_distances, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b8e3b2e4-4888-442d-b7b2-c1a5276cc6c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 将检索到的文档转换为字符串格式\n",
    "context = \"\\n\".join(\n",
    "    [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2d9e034e-c1fd-4ac3-9f4b-e539bd71589b",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 系统和用户提示\n",
    "SYSTEM_PROMPT = \"\"\"\n",
    "Human: You are an AI assistant. You are able to find answers to the questions from the contextual passage snippets provided.\n",
    "\"\"\"\n",
    "USER_PROMPT = f\"\"\"\n",
    "Use the following pieces of information enclosed in <context> tags to provide an answer to the question enclosed in <question> tags.\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a8a1ab9f-116d-4bb0-a14c-fd7645d79c8a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "According to the contextual passage, Milvus stores its data in two places:\n",
      "\n",
      "1. Persistent storage as incremental logs: This includes vector data, scalar data, and collection-specific schema.\n",
      "2. Etcd for metadata: Each Milvus module has its own metadata stored in etcd.\n",
      "\n",
      "Note that metadata is generated within Milvus itself, whereas inserted data are stored in persistent storage as incremental logs.\n"
     ]
    }
   ],
   "source": [
    "# 使用 Ollama 提供的llama3.2 模型，根据提示生成响应\n",
    "from ollama import chat\n",
    "from ollama import ChatResponse\n",
    "\n",
    "response: ChatResponse = chat(\n",
    "    model=\"llama3.2\",\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": USER_PROMPT},\n",
    "    ],\n",
    ")\n",
    "print(response[\"message\"][\"content\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70f12206-854a-43cc-b13d-3eda0da7286f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3.10",
   "language": "python",
   "name": "py310"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
